#!/usr/bin/env python
# -*- coding: utf-8 -*-

#
# This program is free software; you can redistribute it and/or modify  
# it under the terms of the GNU General Public License as published by  
# the Free Software Foundation; either version 2 of the License, or     
# (at your option) any later version.                                   
#                                                                         
# A copy of the license can be found in the license.txt file supplied   
# with this software or at: http://www.gnu.org/copyleft/gpl.html       
#


import time,socket,base64
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream

import celtuce.help.logging as logging
import celtuce.help.url as urlparser
import celtuce.help.codec as codec

from celtuce.helper import getTimestamp
from celtuce.event import Event,ROUTE_GET,ROUTE_DIGEST,ROUTE_DEMAND,OP_RESPONSE,OP_FAILURE
from celtuce.event import OP_INSERT,OP_ERASE,OP_FIND,OP_INSERTUNIQUE,OP_UPDATE,OP_UPSERT
from celtuce.response import Response, RC_TIMEOUT,RC_INTERNAL_ERROR,RC_TABLET_SERVER_NOT_FOUND

from bisect import bisect_left,insort


class   Accessor(object) :
    def __init__(self, meta_url) :
        self._id_base = id(self)
        self._id_serial = 1
        self._id_maximum = 1024 * 1024 * 1024

        self._meta_queue = queue.Queue()
        self._TIMEOUT_WAIT_META_SERVER = 10

        self._data_queue = queue.Queue()
        self._TIMEOUT_WAIT_DATA_SERVER = 10

        self._route_tablet = {}
        self._route_tablet_text = None
        self._route_tablet_digest = None
        self._route_last_update = 0

        self._io_loop = ioloop.IOLoop()
        self._meta_url = meta_url
        self._agent_pool = AgentPool(self._meta_url, self._io_loop)

        self._logger = logging.getLogger('accessor')
        ## __init__()

    def __del__(self) :
        if self._io_loop :
            self.shutdown()
        ## __del__()

    def shutdown(self) :
        self._logger.info('Accessor is shutdowning ...')
        while self._io_loop.running() :
            self._io_loop.stop()
            time.sleep(0.5)
        self._io_loop = None
        ## shutdown()

    ####-----------------------------------------------------------####

    def insert(self, table_name, key, data, timeout_ms=-1) :
        return  self.__doInsertOrUpdate(OP_INSERT, table_name, key, data, timeout_ms)
        ## insert()

    def insertunique(self, table_name, key, data, timeout_ms=-1) :
        return  self.__doInsertOrUpdate(OP_INSERTUNIQUE, table_name, key, data, timeout_ms)
        ## insertunique()

    def __doInsertOrUpdate(self, op, table_name, key, data, timeout_ms=-1) :
        response = Response()

        # get current timestamp
        timestamp_ms = getTimestamp()

        # STEP. lookup tablet-servers
        tablet_servers = self.lookupTabletServer(table_name, key, timeout_ms < 1)
        if not tablet_servers :
            response.status_code = RC_TABLET_SERVER_NOT_FOUND
            response.error = 'No tablet-server for %s(key:%s)' % (table_name, repr(key))
            return response

        # STEP. TODO to check if there is enough tablet-servers 

        # STEP. make INSERT request
        request = Event()
        request.op = op
        request.body = (table_name, key, data, timestamp_ms)
        request.request_id = self.genRequestID()

        # STEP. send request to tablet-servers concurrently, and wait response
        response = self.__callDataServers(request, tablet_servers)

        # STEP. check response; update route-table if necessary

        # STEP. make up response

        return response
        ## __doInsertOrUpdate()

    def erase(self, table_name, key, timeout_ms=-1) :
        response = Response()

        # get current timestamp
        timestamp_ms = getTimestamp()

        # STEP. lookup tablet-servers
        tablet_servers = self.lookupTabletServer(table_name, key, timeout_ms < 1)
        if not tablet_servers :
            response.status_code = RC_TABLET_SERVER_NOT_FOUND
            response.error = 'No tablet-server for %s(key:%s)' % (table_name, repr(key))
            return response

        # STEP. TODO to check if there is enough tablet-servers 

        # STEP. make ERASE request
        request = Event()
        request.op = OP_ERASE
        request.body = (table_name, key, timestamp_ms)
        request.request_id = self.genRequestID()

        # STEP. send request to tablet-servers concurrently, and wait response
        response = self.__callDataServers(request, tablet_servers)

        # STEP. check response; update route-table if necessary

        # STEP. make up response

        return response
        ## erase()

    def find(self, table_name, key, timeout_ms=-1) :
        '''
        @param table_name name of table
        @param key key of record
        @param timeout_ms timeout of find operation (in milliseconds)

        @return 2-tuple (key, value | None)
        '''
        response = Response()

        # get current timestamp
        timestamp_ms = getTimestamp()

        # STEP. lookup tablet-servers
        tablet_servers = self.lookupTabletServer(table_name, key, timeout_ms < 1)
        if not tablet_servers :
            response.status_code = RC_TABLET_SERVER_NOT_FOUND
            response.error = 'No tablet-server for %s(key:%s)' % (table_name, repr(key))
            return response

        # STEP. TODO to check if there is enough tablet-servers 

        # STEP. make FIND request
        request = Event()
        request.op = OP_FIND
        request.body = (table_name, key, timestamp_ms)
        request.request_id = self.genRequestID()

        # STEP. send request to tablet-servers concurrently, and wait response
        response = self.__callDataServers(request, tablet_servers)

        # STEP. check response; update route-table if necessary

        # STEP. merge found data, and return the last one
            
        # STEP. TODO : to do some read-repair

        # STEP. make up response

        return response
        ## find()

    def findall(self, table_name, key, timeout_ms=-1) :
        '''
        @param table_name name of table
        @param key key of record
        @param timeout_ms timeout of find operation (in milliseconds)

        @return 3-tuple (key, count-of-values, list-of-values | None)
        '''
        response = Response()

        # get current timestamp
        timestamp_ms = getTimestamp()

        # STEP. lookup tablet-servers
        tablet_servers = self.lookupTabletServer(table_name, key, timeout_ms < 1)
        if not tablet_servers :
            response.status_code = RC_TABLET_SERVER_NOT_FOUND
            response.error = 'No tablet-server for %s(key:%s)' % (table_name, repr(key))
            return response

        # STEP. TODO to check if there is enough tablet-servers 

        # STEP. make FIND request
        request = Event()
        request.op = OP_FIND
        request.body = (table_name, key, timestamp_ms)
        request.request_id = self.genRequestID()

        # STEP. send request to tablet-servers concurrently, and wait response
        replies = self.__callDataServers(request, tablet_servers)

        # STEP. check response; update route-table if necessary
        unique_dataset = []
        i = 0
        for reply_from,reply in replies :
            i += 1
            if reply.isOk() :
                self._logger.debug('FINDALL #%d reply (from %s) body : %s', i, reply_from, reply.body)   
                self.__mergeRowSet(unique_dataset, reply.body)
            else :
                self._logger.debug('FINDALL #%d reply (from %s) error : %s', i, reply_from, reply.error)   

        # STEP. merge found data
        #       make sure all data is not delete
        index_delete = len(unique_dataset) - 1
        while index_delete >= 0 :
            _op = unique_dataset[ index_delete ][ 1 ]
            if _op != 0 :
                # not insert
                break
            index_delete -= 1
        if index_delete >= 0 :
            unique_dataset = unique_dataset[ index_delete + 1 : ]
            
        dataset = []
        for ts,op,data in unique_dataset :
            dataset.append( data )

        # STEP. TODO : to do some read-repair

        # STEP. make up response
        response.body = [key, len(dataset), dataset]

        return response
        ## findall()


    def __mergeRowSet(self, unique_dataset, record) :
        if not record : 
            return

        try :
            key,op,ts,data = record
        except :
            return

        row = (ts, op, data)
        found = False

        index = bisect_left(unique_dataset, row)
        if index < len(unique_dataset) :
            _row = unique_dataset[ index ]
            _ts,_op = _row[0:2]
            if _ts == ts and _op == op :
                found = True

        if not found :
            insort(unique_dataset, row)
        ## __mergeRowSet()

    ####-----------------------------------------------------------####

    def needUpdateRouteTable(self) :
        '''
        @return True on success, False on failure
        '''
        request = Event()
        request.request_id = self.genRequestID()
        request.op = ROUTE_DIGEST

        response = self.__callMetaServer( request )
        if not response :
            self._logger.error('fail to check route table digest.')
            return False
        if not response.isOk() :
            self._logger.error('fail to check route table digest -- %s', response.error)
            return False

        route_tablet_digest = response.body
        if route_tablet_digest == self._route_tablet_digest :
            # the same digest, route table is updated
            return False

        return True
        ## needUpdateRouteTable()

    def updateRouteTable(self) :
        '''
        @return True on success, False on failure
        '''
        request = Event()
        request.request_id = self.genRequestID()
        request.op = ROUTE_GET

        response = self.__callMetaServer( request )
        if not response :
            self._logger.error('fail to update route table.')
            return False
        if not response.isOk() :
            self._logger.error('fail to update route table -- %s', response.error)
            return False

        # TODO : parse and load route table
        route_data = response.body
        if isinstance(route_data, dict) :
            route_data = [ route_data[i] for i in range(len(route_data)) ]
        if route_data[1] != self._route_tablet_digest :
            self._route_tablet_text,self._route_tablet_digest = route_data
            #self._logger.info('to update route table with %s', repr(route_data))

            self._route_tablet = codec.decodeRouteTablet( self._route_tablet_text )
            self._logger.info('ROUTE tablet : %s', repr(self._route_tablet))

        # last check consist of route-table
        self._route_last_update = getTimestamp()

        return True
        ## updateRouteTable()


    def demandRoute(self, table_name, key) :
        '''
        @return True on success, False on failure
        '''
        request = Event()
        request.request_id = self.genRequestID()
        request.op = ROUTE_DEMAND
        request.body = (table_name, key)

        response = self.__callMetaServer( request )
        if not response :
            self._logger.error('fail to demand route.')
            return False
        if not response.isOk() :
            self._logger.error('fail to demand route -- %s', response.error)
            return False

        result = response.body

        return True
        ## demandRoute()

    ####-----------------------------------------------------------####

    def lookupTabletServer(self, table_name, key, blocked=True) :
        '''lookup tablet-server

        @param table_name name of table
        @param key key of record
        @param blocked whether blocked if tablet-server not found

        @return list of tablet-server charging (table_name, key); \
                None if not found in non-blocking mode
        '''
        count = 0
        while True :
            tablet_servers = self.__lookupTabletServer(table_name, key) 
            if tablet_servers :
                if len(tablet_servers) < self.__getExpectReplicaNumber(table_name) :
                    # if more tablet-servers are expected, update table_name periodly 
                    if (getTimestamp() - self._route_last_update) > 60000 :
                        self.updateRouteTable()

                return tablet_servers

            # got nothing
            if (not blocked) and (count > 0) :
                return None

            if count == 0 :
                count += 1
            else :
                time.sleep(1.0)
            
            # got nothing 
            # try to alloc route-item
            if not self.demandRoute(table_name, key) :
                continue
            
            # try to update route-table
            if not self.updateRouteTable() :
                continue

            # end of while loop
        ## lookupTabletServer()

    def __lookupTabletServer(self, table_name, key) :
        if table_name not in self._route_tablet :
            return None

        tablets = self._route_tablet[ table_name ]
        index = bisect_left(tablets, (key,key,None,None))
        if index > 0 :
            index -= 1

        key_start,key_stop,ts_list = tablets[ index ]
        if (key_start < key) and (key <= key_stop) :
            return ts_list

        return None
        ## __lookupTabletServer()

    ####-----------------------------------------------------------####

    def genRequestID(self) :
        # TODO: lock
        serial = self._id_serial

        self._id_serial += 1
        if self._id_serial >= self._id_maximum :
            self._id_serial = 1 

        id = '.'.join((str(self._id_base), str(serial)))
        return id
        ## genRequestID()

    def __getExpectReplicaNumber(self, table_name) :
        return  3
        ## __getExpectReplicaNumber()

    ####-----------------------------------------------------------####


    def __callMetaServer(self, request) :

        # send request
        response = self._agent_pool.callMetaServer( request, self._TIMEOUT_WAIT_META_SERVER )

        return response
        ## __callMetaServer()


    def __callDataServers(self, request, tablet_servers, timeout_ms=0) :
        '''
        @param request request to be sent
        @param tablet_servers name list of tablet-server
        @param timeout_ms timeout when calling tablet-server

        @return list of (tablet_server_address, response) tuples
        '''

        if timeout_ms > 0 :
            timeout_s = timeout_ms / 1000.0
        else :
            timeout_s = self._TIMEOUT_WAIT_DATA_SERVER

        # send request
        response = self._agent_pool.callDataServer(request, tablet_servers, timeout_s)

        return response
        ## __callDataServers()

    ## class Accessor()

#######################################################################

try :
    from hashlib import md5 as MD5
except :
    from md5 import md5 as MD5
import celtuce.help.url as urlparser


class   AgentPool(object) :

    def __init__(self, meta_url, io_loop) :
        self._io_loop = io_loop

        # such as tcp://host:port
        self._meta_url = meta_url
        parser = urlparser.URL(meta_url)
        self._meta_protocol = parser.protocol
        self._meta_host = parser.host
        self._meta_port = parser.port

        self._END_TAG = "\r\n\r\n"

        self._agents = {}
        self._requests = {}

        self._id_base = id(self)
        self._id_serial = 1 
        self._id_maximum = 1024 * 1024 * 1024

        self._when_to_update_route_table = 1.0

        self._logger = logging.getLogger('agentpool')
        ## __init__()

    def callDataServer(self, request, tablet_servers, timeout_s) :
        # STEP. lookup agent for this request
        agents = self.lookupAgent( tablet_servers )
        if not agents :
            self._logger.error('fail to lookup agent for tablet-servers %s', tablet_servers)

            response = Response(RC_INTERNAL_ERROR)
            response.error = 'fail to lookup agent for tablet-servers %s' % tablet_servers
            return response

        # STEP. put request into request-pool
        timestamp = getTimestamp()
        result = {}
        for agent in agents :
            result[ agent ] = None
        self._requests[ (request, timestamp) ] = result

        # STEP. put request into queue of the agent
        for agent in agents :
            agent.sendRequest( request, timestamp, timeout_s )

        # STEP. run request
        self._timeout_called = False
        deadline = time.time() + timeout_s
        self._timeout_handle = self._io_loop.add_timeout( deadline, self.__onTimeoutCall )

        self._data_response = None
        self._io_loop.start()

        if self._timeout_called :
            response = Response(RC_TIMEOUT)
            response.error = 'Timeout to wait reply for request "%s"' % str(request)
            return response

        del self._requests[ (request, timestamp) ]

        return self._data_response
        ## handleClientRequest()


    def feedResponse(self, request, timestamp, agent, response) :
        id = (request, timestamp)
        if id not in self._requests :
            # just discard some out-of-date response
            self._logger.error('response for (%s, %s) is out of date', str(request), str(timestamp))
            return

        result_set = self._requests[ id ]
        if agent not in result_set :
            self._logger.error('unknown agent for (%s, %s)', str(request), str(timestamp))
            return

        result_set[ agent ] = response

        # check that all response is returned
        want_more_response = False
        for _agent,_response in result_set.iteritems() :
            if _response is None :
                want_more_response = True
                break
        if want_more_response :
            return

        # we got all response
        self._io_loop.remove_timeout( self._timeout_handle )
        self._io_loop.stop()
        
        # merge all responses, and return
        done_set = []
        fail_set = []
        for _agent,_response in result_set.iteritems() :
            r = _response.body
            if r.isOk() :
                done_set.append( (_agent, r) )
            else :
                fail_set.append( (_agent, r) )

        # TODO : to find the optimal result
        if done_set :
            # success
            response_packet = self.__makeSuccessResponsePacket(done_set, request)
        else :
            response_packet = self.__makeFailureResponsePacket(request)
            # failure

        # handle error response
        if fail_set :
            self.__handleErrorResponses(done_set, fail_set)

        self._data_response = response_packet
        ## feedResponse()

    def __makeSuccessResponsePacket(self, done_set, request) :
        response = Response()

        if request.op not in (OP_FIND,) :
            return response

        # merge result
        if len(done_set) == 1 :
            for _agent,_response in done_set :
                #self._logger.debug('response body : %s', _response.body)
                if isinstance(_response.body, dict) :
                    _key,_op,_ts,_data = [ _response.body[i] for i in range(len(_response.body)) ]
                else :
                    _key,_op,_ts,_data = _response.body
                response.body = (_key, _data)
                break
        else :
            result_set = []
            for _agent,_response in done_set :
                if _response.body :
                    #self._logger.debug('response body : %s', _response.body)
                    if isinstance(_response.body, dict) :
                        _key,_op,_ts,_data = [ _response.body[i] for i in range(len(_response.body)) ]
                    else :
                        _key,_op,_ts,_data = _response.body
                    insort(result_set, (_ts, _op, _key, _data))
            last = result_set[ -1 ]
            _ts,_op,_key,_data = result_set[ -1 ]
            response.body = (_key, _data)

        return response
        ## __makeSuccessResponsePacket()

    def __makeFailureResponsePacket(self, request) :
        response = Response(RC_INTERNAL_ERROR)

        return response
        ## __makeFailureResponsePacket()

    def __handleErrorResponses(self, done_set, fail_set) :
        for _agent,r in fail_set :
            self._logger.error('Error Response : %s', r.error)
        ## __handleErrorResponses()

    ####-----------------------------------------------------------####

    def __onTimeoutCall(self) :
        self._io_loop.stop()
        self._timeout_called = True
        ## __onTimeoutCall()

    def callMetaServer(self, request, timeout_s) :
        self._logger.debug('access Meta-Server %s:%s ...', self._meta_host, self._meta_port)

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._meta_host, self._meta_port) )
        except :
            self._logger.excpt()
            
            response = Response(RC_INTERNAL_ERROR)
            response.error = 'fail to connect to meta-server (tcp://%s:%d)' % (self._meta_host, self._meta_port)
            return response

        stream = iostream.IOStream(s, io_loop=self._io_loop)
        MetaConnection(stream, (self._meta_host, self._meta_port), request, self.onMetaServerCallback)

        self._timeout_called = False
        deadline = time.time() + timeout_s
        self._timeout_handle = self._io_loop.add_timeout( deadline, self.__onTimeoutCall )

        self._meta_response = None
        self._io_loop.start()

        if self._timeout_called :
            response = Response(RC_TIMEOUT)
            response.error = 'Timeout to wait reply for request "%s"' % str(request)
            return response

        return self._meta_response
        ## callMetaServer()

    def onMetaServerCallback(self, response) :
        self._io_loop.remove_timeout( self._timeout_handle )
        self._io_loop.stop()

        if response : 
            self._meta_response = Response()
            self._meta_response.body = response.body
        else :
            self._meta_response = Response(RC_INTERNAL_ERROR)
            self._meta_response.error = "Empty response."
        ## onMetaServerCallback()        

    ####-----------------------------------------------------------####

    def lookupAgent(self, tablet_servers) :
        agents = []
        for ts in tablet_servers :
            if ts not in self._agents :
                agent = Agent(self, ts, self._io_loop)
                self._agents[ ts ] = agent
            agents.append( self._agents[ ts ] )
       
        return agents
        ## lookupAgent()


    ####-----------------------------------------------------------####

    def genRequestID(self) :
        # TODO: lock
        serial = self._id_serial

        self._id_serial += 1
        if self._id_serial >= self._id_maximum :
            self._id_serial = 1

        id = '.'.join((str(self._id_base), str(serial)))
        return id
        ## genRequestID()


    ## class AgentPool

class   Agent(object) :
    def __init__(self, agent_pool, tablet_server_address, io_loop) :
        self._ts_address = tablet_server_address
        self._agent_pool = agent_pool
        self._io_loop = io_loop

        self._ongoing = False
        self._keep_alive = True

        self._stream = None

        self._END_TAG = "\r\n\r\n"

        self._client_connection = None
        self._logger = logging.getLogger('agent')
        ## __init__()

    def sendRequest(self, request, timestamp, timeout_s) :
        self._tag_ongoing = (request, timestamp)
        self._wait_response_for = request.request_id
        request_str = codec.encode( request )
        self._request_ongoing = request_str

        self._timeout = time.time() + timeout_s

        self._ongoing = False

        self.__handleRequest()
        ## sendRequest()

    def __handleRequest(self) :
        if self._ongoing and self._stream :
            return

        self._ongoing = True

        ok = self.__connect()
        if not ok :
            return

        self._stream.write( ''.join((self._request_ongoing, self._END_TAG)) )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __handleRequest()    

    def __onResponse(self, response_str) :

        response = codec.decode( response_str )
        if response.response_to == self._wait_response_for :
            # put response to request pool
            _request, _timestamp = self._tag_ongoing
            self._agent_pool.feedResponse(_request, _timestamp, self, response)

            # handle any queued request
            self._ongoing = False
            self._tag_ongoing = None
            self._request_ongoing = None
        else :
            # TODO : ignore other response and request
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    def __connect(self) :
        if self._stream :
            return True

        try :
            self._logger.debug('try to connect to %s ...', self._ts_address)
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( self._ts_address )
            self._logger.info('connect to %s done.', self._ts_address)
        except :
            s.close()
            # fail to connect
            self._logger.excpt()

            time_now = time.time()

            if time_now < self._timeout :
                deadline = time_now + 1.0
                self._io_loop.add_timeout(deadline, self.__tryToConnect)
            else :
                response = Response(RC_TIMEOUT)
                response.error = 'timeout when trying to connect to %s' % self._ts_address

                _request, _timestamp = self._tag_ongoing
                self._agent_pool.feedResponse(_request, _timestamp, self, response)
            return False

        self._stream = iostream.IOStream( s, self._io_loop )
        self._stream.set_close_callback( self.__onDisconnect )
        return True
        ## __connect()

    def __disconnect(self, force=False) :
        if self._keep_alive and not force :
            return

        self._stream.close()
        self._stream = None
        ## __disconnect()

    def __onDisconnect(self) :
        self._logger.error('ERROR : connection to %s broke!', self._ts_address)
        self.__disconnect( force=True )

        deadline = time.time() + 1.0
        self._io_loop.add_timeout(deadline, self.__tryToConnect)
        ## __onDisconnect()

    def __tryToConnect(self) :
        if self._stream :
            return

        self.__handleRequest()
        ## __tryToConnect()

    ## class Agent

class   MetaConnection(object) :
    def __init__(self, stream, address, request, response_callback, end_callback=None) :
        self._stream = stream
        self._address = address
        self._response_callback = response_callback
        self._end_callback = end_callback

        self._keep_alive = False

        self._END_TAG = "\r\n\r\n"
        self._request_finished = False

        self._logger = logging.getLogger('metaconnect')
        self._logger.debug('create MetaConnection for %s', self._address)

        self._wait_response_for = request.request_id
        request_str = codec.encode(request)
        self._stream.write( ''.join((request_str, self._END_TAG)) )
        self._stream.read_until(self._END_TAG, self.onResponse)
        ## __init__()

    def __del__(self) :
        #if self._stream.closed() :
        #    self._logger.debug('stream is closed when destroying')
        #else :
        #    self._logger.debug('stream is opened when destroying')
        self._logger.debug('destroy MetaConnection for %s', self._address)
        pass
        ## __del__()

    def onResponse(self, response_str) :
        response = codec.decode( response_str )
        #self._logger.debug('response_to=%s wait_response=%s', response.response_to, self._wait_response_for)
        if response.response_to == self._wait_response_for :
            self._response_callback( response )

            self._stream.close()
            if self._end_callback : self._end_callback()
        else :
            # TODO : ignore other response and request
            self._stream.read_until(self._END_TAG, self.onResponse)
        ## onResponse()


    ## class MetaConnection

#######################################################################



